from typing_extensions import TypedDict
from typing import Annotated, Literal
from langgraph.graph.message import add_messages
from langchain_core.tools import tool
from langgraph.prebuilt import ToolNode
from langchain_community.chat_models import ChatTongyi
from langgraph.graph import StateGraph, END
# from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.messages import HumanMessage, AIMessage

from dotenv import load_dotenv
load_dotenv(verbose=True)

# Add messages essentially does this with more robust handling
# def add_messages(left: list, right: list):
#     return left + right


class State(TypedDict):
    messages: Annotated[list, add_messages]


@tool
def searchWeb(query: str):
    """Call to surf the web."""
    # This is a placeholder for the actual implementation
    return ["The weather is cloudy with a chance of meatballs."]


tools = [searchWeb]
tool_node = ToolNode(tools)

# We will set streaming=True so that we can stream tokens
# See the streaming section for more information on this.
model = ChatTongyi(
    model="qwen-turbo",
    temperature=0,
)

model = model.bind_tools(tools)


# 定义方法判断是否继续
def should_continue(state: State) -> Literal["continue", "end"]:
    last_message = state["messages"][-1]
    # If there is no function call, then we finish
    if not last_message.tool_calls:
        return "end"
    # Otherwise if there is, we continue
    else:
        return "continue"


def call_model(state: State):
    response = model.invoke(state["messages"])
    # We return a list, because this will get added to the existing list
    return {"messages": response}


# Define a new graph
workflow = StateGraph(State)

# Define the two nodes we will cycle between
workflow.add_node("agent", call_model)
workflow.add_node("action", tool_node)

# Set the entrypoint as `agent`
# This means that this node is the first one called
workflow.set_entry_point("agent")

# We now add a conditional edge
workflow.add_conditional_edges(
    # First, we define the start node. We use `agent`.
    # This means these are the edges taken after the `agent` node is called.
    "agent",
    # Next, we pass in the function that will determine which node is called next.
    should_continue,
    # Finally we pass in a mapping.
    # The keys are strings, and the values are other nodes.
    # END is a special node marking that the graph should finish.
    # What will happen is we will call `should_continue`, and then the output of that
    # will be matched against the keys in this mapping.
    # Based on which one it matches, that node will then be called.
    {
        # If `tools`, then we call the tool node.
        "continue": "action",
        # Otherwise we finish.
        "end": END,
    },
)

# We now add a normal edge from `tools` to `agent`.
# This means that after `tools` is called, `agent` node is called next.
workflow.add_edge("action", "agent")

#memory = SqliteSaver.from_conn_string(":memory:")
memory = MemorySaver()
print(memory)
# Finally, we compile it!
# This compiles it into a LangChain Runnable,
# meaning you can use it as you would any other runnable
app = workflow.compile(checkpointer=memory)


def main():
    nowStats = {}
    # workflow = StateGraph(State)
    # app = workflow.compile(checkpointer=memory)

    config = {"configurable": {"thread_id": "2"}}
    input_message = HumanMessage(content="hi! I'm bob")
    for event in app.stream({"messages": [input_message]},
                            config,
                            stream_mode="values"):
        event["messages"][-1].pretty_print()

    #获取状态
    nowStats = app.get_state(config).values
    print(nowStats)
    #获取下一步状态
    nowStats = app.get_state(config).next
    print(nowStats)

    #执行工具
    config = {"configurable": {"thread_id": "2"}}
    input_message = HumanMessage(content="what is the weather in sf currently")
    for event in app.stream({"messages": [input_message]},
                            config,
                            stream_mode="values"):
        event["messages"][-1].pretty_print()

    print(
        "************************************thread_id:4*****************************************************"
    )
    #重新定义工作流，注意增加了interrupt_before参数 这意味着在采取任何操作之前我们会暂停，人类可以在这个时候插手AI工作
    app_w_interrupt = workflow.compile(checkpointer=memory,
                                       interrupt_before=["action"])
    #请求工具
    config = {"configurable": {"thread_id": "4"}}
    input_message = HumanMessage(content="what is the weather in sf currently")
    for event in app_w_interrupt.stream({"messages": [input_message]},
                                        config,
                                        stream_mode="values"):
        event["messages"][-1].pretty_print()
    #获取当前节点状态以及下一个节点状态
    current_values = app_w_interrupt.get_state(config)
    print("====当前节点状态====>\n")
    print(current_values)
    print("===下一个节点状态====>\n")
    print(current_values.next)
    print("===当前的工具输入内容====>\n")
    print(current_values.values["messages"][-1].tool_calls)
    #在暂停的时候将工具输入从Agent给的指令修改为人类的
    current_values.values["messages"][-1].tool_calls[0]["args"][
        "query"] = "洛杉矶今天的天气如何？"
    #调用update_state方法，将工具输入的内容修改为人类的输入
    app_w_interrupt.update_state(config, current_values.values)
    #查看当前工作流中工具的输入值
    print("====修改后的工具输入====>\n")
    print(app_w_interrupt.get_state(config).values)
    #再次启动Agent,并从中断的节点处恢复运作，注意设置为None即是从当前开始
    print("====再次启动agent====>\n")
    for event in app_w_interrupt.stream(None, config):
        for v in event.values():
            print(v)
    #回溯整个工作流
    print("---------------回溯工作流:-------->>>>>>>>>>>>>>>>>>>>>>>>>>>>")
    for state in app_w_interrupt.get_state_history(config):
        print(state)
        print("--")
        #第二条消息
        if len(state.values["messages"]) == 2:
            to_replay = state
    print("====回溯的起点节点：====>\n")
    print(to_replay.values)
    #重播工作流
    print("---------------开始重播工作流:-------->>>>>>>>>>>>>>>>>>>>>>>>>>>>")
    for event in app_w_interrupt.stream(None, to_replay.config):
        for v in event.values():
            print(v)

    #不仅是重播，甚至可以改写过去的流向，比如提前终止工作流
    print("---------------开始分支工作流:-------->>>>>>>>>>>>>>>>>>>>>>>>>>>>")
    branch_config = app_w_interrupt.update_state(
        to_replay.config,
        {
            "messages": [
                AIMessage(content="All done here!",
                          id=to_replay.values["messages"][-1].id)
            ]
        },
    )
    branch_state = app_w_interrupt.get_state(branch_config)
    print("====分支前的节点：====>\n")
    print(branch_state.values)
    print("====运行分支后的节点：====>\n")
    print(branch_state.next)


if __name__ == "__main__":
    main()
